
#include "KdTree.h"

#include <float.h>
#include <math.h>
#include <iomanip>
#include <iostream>
using namespace std;

KdTree::KdTree(INT_PAIR dim, int num_procs) : dim(dim), num_procs(num_procs),
  proc_count(0) {
  root_node = new KdTreeNode();
  tree_depth = (int)(log2(num_procs)+0.5);

  INT_PAIR min(0, 0), max(dim[0], dim[1]); 
  buildKdTree(root_node, 0, min, max);
  generateProcAllocTable();
}

KdTree::~KdTree() {
  delete root_node;
}

void KdTree::buildKdTree(KdTreeNode* node, int depth, INT_PAIR min, INT_PAIR max) {
  node->min = min;
  node->max = max;

  node->depth = depth; 

  if (depth==tree_depth) {
   proc_node_map[proc_count++] = node;   
   return;
  }

  int axes = node->axes = depth%2;
  int split_plane_pos = 0;

  // Use spatial median to determine the split plane position
  if (axes==0)
    split_plane_pos = (max[0]+min[0])/2;
  else if (axes==1)
    split_plane_pos = (max[1]+min[1])/2;

  // node->split_plane_pos = split_plane_pos;

  // Create two partitions
  INT_PAIR min1(min[0], min[1]);
  INT_PAIR max2(max[0], max[1]);

  INT_PAIR max1, min2;
    
  if (axes==0) { // x-axis
    max1 = INT_PAIR(split_plane_pos, max2[1]);
    min2 = INT_PAIR(split_plane_pos, min1[1]);
  }
  else if (axes==1) { // y-axis   	  	 	    	  	 	 
    max1 = INT_PAIR(max2[0], split_plane_pos);
    min2 = INT_PAIR(min1[0], split_plane_pos);
  }

  KdTreeNode* leftNode = node->left = new KdTreeNode();
  KdTreeNode* rightNode = node->right = new KdTreeNode();

  buildKdTree(leftNode, depth+1, min1, max1);
  buildKdTree(rightNode, depth+1, min2, max2);                  
}

void KdTree::generateProcAllocTable() {
  // printf("<<<Unique elements inserted into the list-0>>>\n");

  // Populate the two processor maps
  for (int i=0; i<proc_count; i++) {
    KdTreeNode *node = proc_node_map[i];
    INT_PAIR min = node->min;
    INT_PAIR max = node->max;

    INT_PAIR range_x(min[0], max[0]);
    INT_PAIR range_y(min[1], max[1]);

    if (proc_x_map.find(range_x)==proc_x_map.end())
      proc_x_map.insert(range_x);
    
    if (proc_y_map.find(range_y)==proc_y_map.end())
      proc_y_map.insert(range_y);
  }

  // printf("<<<Unique elements inserted into the list-1>>>\n");

  // Re-size the processor allocation table
  proc_alloc_table.resize(proc_y_map.size());
  for (int i=0 ; i<proc_y_map.size(); i++)
    proc_alloc_table[i].resize(proc_x_map.size());

  // Populate the processor allocation table
  set<INT_PAIR>::iterator range_iterator;
  int range_index_x, range_index_y;  
  for (int i=0; i<proc_count; i++) {
    KdTreeNode *node = proc_node_map[i];
    INT_PAIR min = node->min;
    INT_PAIR max = node->max;

    INT_PAIR range_x(min[0], max[0]);
    INT_PAIR range_y(min[1], max[1]);

    if ((range_iterator=proc_x_map.find(range_x))!=\
	proc_x_map.end()) {
      range_index_x = distance(proc_x_map.begin(), range_iterator);
    }
    else {
      range_index_x = -1;
    }
    
    if ((range_iterator=proc_y_map.find(range_y))!=\
	proc_y_map.end()) {
      range_index_y = distance(proc_y_map.begin(), range_iterator);
    }
    else {
      range_index_y = -1;
    }

    if (range_index_x!=-1&&range_index_y!=-1)
      proc_alloc_table[range_index_y][range_index_x] = i;
  }
}

const std::vector< std::vector<int> >& KdTree::getProcAllocTable() const {
  return proc_alloc_table;
}

int KdTree::getNumBlocksX() const {
  return proc_x_map.size();
}

int KdTree::getNumBlocksY() const {
  return proc_y_map.size();
}

const KdTreeNode* KdTree::getProcNode(int proc) {
  if (proc_node_map.find(proc)!=proc_node_map.end())
    return proc_node_map[proc];
  else
    return NULL;
}

